package com.lucifer.hawkeye.ai.config;

import jakarta.annotation.Resource;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.document.MetadataMode;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.openai.OpenAiEmbeddingModel;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.ai.zhipuai.ZhiPuAiChatModel;
import org.springframework.ai.zhipuai.ZhiPuAiEmbeddingModel;
import org.springframework.ai.zhipuai.ZhiPuAiEmbeddingOptions;
import org.springframework.ai.zhipuai.api.ZhiPuAiApi;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.jdbc.DataSourceBuilder;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.JdbcTemplate;
import javax.sql.DataSource;
import java.util.Objects;

/**
 * @author lucifer
 * @date 2025/3/29 11:05
 */
@Configuration
public class HawkeyeAiConfig {


    @Value("${spring.ai.custom.model}")
    private String model;

    @Value("${spring.ai.custom.queryTransformer}")
    public boolean queryTransformer;

    @Value("${spring.ai.custom.knowledge}")
    public boolean knowledge;

    @Value("${spring.ai.zhipuai.api-key}")
    private String zhipuaiApiKey;

    @Value("${spring.ai.openai.api-key}")
    private String openaiApiKey;


    @Value("${spring.ai.openai.base-url}")
    private String openaiBaseUrl;

    @Value("${spring.datasource.pgvector.driverClassName}")
    private String driverClassName;

    @Value("${spring.datasource.pgvector.url}")
    private String url;

    @Value("${spring.datasource.pgvector.username}")
    private String username;

    @Value("${spring.datasource.pgvector.password}")
    private String password;



    @Bean(name = "zhiPuEmbeddingModel")
    public EmbeddingModel zhiPuEmbeddingModel() {
        ZhiPuAiApi zhiPuAiApi = new ZhiPuAiApi(zhipuaiApiKey);
        return new ZhiPuAiEmbeddingModel(zhiPuAiApi, MetadataMode.EMBED,
                ZhiPuAiEmbeddingOptions.builder()
                        .model("embedding-2")
                        .dimensions(1024)
                        .build());
    }


   @Resource(name = "openAiEmbeddingModel")
   private OpenAiEmbeddingModel openAiEmbeddingModel;


    @Bean(name = "pgJdbcTemplate")
    public JdbcTemplate pgJdbcTemplate() {
        DataSource dataSource = DataSourceBuilder.create().url(url)
                .username(username)
                .password(password)
                .driverClassName(driverClassName)
                .build();
        return new JdbcTemplate(dataSource);
    }


    @Bean(name = "docVectorStore")
    public VectorStore docVectorStore() {
        PgVectorStore pgVectorStore = PgVectorStore.builder(pgJdbcTemplate(), embeddingModel())
                .distanceType(PgVectorStore.PgDistanceType.COSINE_DISTANCE)       // Optional: defaults to COSINE_DISTANCE
                .indexType(PgVectorStore.PgIndexType.HNSW)                     // Optional: defaults to HNSW
                .initializeSchema(false)              // Optional: defaults to false
                .schemaName("public")                // Optional: defaults to "public"
                .vectorTableName("doc_store")     // Optional: defaults to "vector_store"
                .dimensions(1024)
                .maxDocumentBatchSize(10000)         // Optional: defaults to 10000
                .build();
        return pgVectorStore;
    }

    @Bean(name = "chatMemoryVectorStore")
    public PgVectorStore chatMemoryVectorStore() {
        PgVectorStore pgVectorStore = PgVectorStore.builder(pgJdbcTemplate(), embeddingModel())
                .distanceType(PgVectorStore.PgDistanceType.COSINE_DISTANCE)       // Optional: defaults to COSINE_DISTANCE
                .indexType(PgVectorStore.PgIndexType.HNSW)                     // Optional: defaults to HNSW
                .initializeSchema(false)              // Optional: defaults to false
                .schemaName("public")                // Optional: defaults to "public"
                .vectorTableName("chat_memory")     // Optional: defaults to "vector_store"
                .maxDocumentBatchSize(10000)         // Optional: defaults to 10000
                .build();
        return pgVectorStore;
    }

    @Bean(name = "nl2sqlVectorStore")
    public PgVectorStore nl2sqlVectorStore() {
        PgVectorStore pgVectorStore = PgVectorStore.builder(pgJdbcTemplate(), embeddingModel())
                .distanceType(PgVectorStore.PgDistanceType.COSINE_DISTANCE)       // Optional: defaults to COSINE_DISTANCE
                .indexType(PgVectorStore.PgIndexType.HNSW)                     // Optional: defaults to HNSW
                .initializeSchema(false)              // Optional: defaults to false
                .schemaName("public")                // Optional: defaults to "public"
                .vectorTableName("nl2sql_store")     // Optional: defaults to "nl2sql_store"
                .maxDocumentBatchSize(10000)         // Optional: defaults to 10000
                .build();
        return pgVectorStore;
    }


    @Bean
    public TokenTextSplitter tokenTextSplitter() {
        return new TokenTextSplitter();
    }

    private EmbeddingModel embeddingModel(){
        if(Objects.equals("zhipuai",model)){
            return zhiPuEmbeddingModel();
        }else if (Objects.equals("openai",model)){
            return openAiEmbeddingModel;
        }else {
            return openAiEmbeddingModel;
        }
    }

    @Resource
    private OpenAiChatModel openAiChatModel;

    @Resource
    private ZhiPuAiChatModel zhiPuAiChatModel;


    public ChatModel chatModel() {
        if (Objects.equals("zhipuai", model)) {
            return zhiPuAiChatModel;
        } else if (Objects.equals("openai", model)) {
            return openAiChatModel;
        }  else {
            return openAiChatModel;
        }
    }

}
